from .csv_group import CSVGroupingEngine
from .dir_group import DirectoryGroupingEngine
from models import Project
from typing import Set, Union
import os


_engines = {
    "csv": CSVGroupingEngine,
    "directory": DirectoryGroupingEngine
}


def create_group_engine(project: Project, selected_categories: Set[str],
                        padding: bool = False, resize: bool = False,
                        csv_mode: bool = False) -> Union[CSVGroupingEngine, DirectoryGroupingEngine]:
    global _engines
    mode = "csv" if csv_mode else "directory"
    target_location = os.path.join(project.location, "exports")
    if not os.path.exists(target_location):
        os.makedirs(target_location)
    return _engines[mode](target_location,
                          project.export_file,
                          project.export_sample,
                          project.export_validate,
                          project.sample_long_side if resize else 0,
                          padding,
                          project.sample_type,
                          selected_categories,
                          project.training_groups,
                          project.training_rate,
                          project.validation_rate)
